import os
from tqdm import tqdm
import sys
import os
import json
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "large_models"))
from qwen_vl import QwenVLModel
from deepseek import DeepseekModel
from PIL import Image
from markdown_popup import show_markdown_popup

def check_image_dimensions(image_dir):
    """
    检查图片目录中所有图片的尺寸是否符合要求（大于1000x1000且为正方形）。
    
    参数:
        image_dir (str): 图片目录路径。
    
    返回:
        list: 符合尺寸要求的图片文件名列表。
    
    异常:
        如果发现不符合尺寸要求的图片，会打印错误信息并退出程序。
    """
    image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.jpg', '.jpeg'))]
    valid_images = []
    invalid_images = []
    for image_file in tqdm(image_files, desc="检查图片尺寸"):
        image_path = os.path.join(image_dir, image_file)
        try:
            with Image.open(image_path) as img:
                width, height = img.size
                if width >= 1000 and height >= 1000 and width == height:
                    valid_images.append(image_file)
                else:
                    invalid_images.append((image_file, f"{width}x{height}"))
        except Exception as e:
            print(f"无法检查图片 {image_file} 的尺寸: {e}")
    
    if invalid_images:
        print("以下图片尺寸不符合要求（需为大于1000x1000的正方形）：")
        for image_file, size in invalid_images:
            print(f"图片名称: {image_file}, 尺寸: {size}")
        sys.exit(1)
    
    return valid_images

def get_image_descriptions(product_name, image_dir):
    """
    获取图片的描述信息。
    
    参数:
        product_name (str): 产品名称，用于生成描述提示。
        image_dir (str): 图片目录路径。
    
    返回:
        tuple: 包含两个列表的元组：
            - 第一个列表是图片文件名和描述的元组列表。
            - 第二个列表是描述中包含中文的图片文件名列表。
    """
    descriptions = []
    chinese_images = []
    image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.jpg', '.jpeg'))]
    model = QwenVLModel(model='qwen-vl-max-2025-04-08')
    
    for image_file in tqdm(image_files, desc="获取图片描述..."):
        image_path = os.path.join(image_dir, image_file)
        prompt = f"这是关于{product_name}的产品图片之一，请详细描述图片。描述包括且不限于：包含的物品描述、图片风格描述、优点、缺点等。以markdown格式输出结果。"
        try:
            response = model.chat(image_path, prompt)
            descriptions.append((image_file, response))
            # 检查描述是否包含中文
            # if any('\u4e00' <= char <= '\u9fff' for char in response):
            #     chinese_images.append(image_file)
        except Exception as e:
            print(f"Error processing {image_file}: {e}")
    
    return descriptions, chinese_images

def analyze_descriptions(product_name, descriptions):
    """
    分析图片描述并提出补充建议。
    
    参数:
        product_name (str): 产品名称，用于生成分析提示。
        descriptions (list): 图片文件名和描述的元组列表。
    
    返回:
        str: 分析结果和建议的文本。
    """
    # prompt = f"下面是产品{product_name}一套电商图片的描述。针对我们的目标平台，一套产品图片最少是5张，最多10张，必须具有的图片是：首图，尺寸图（至少一个），场景图（至少一个）。其它可选图片包括但不限于：用例图（可以多个），细节图（可以多个），模特图，侧面图，内部图。请根据描述，对图片进行分类，然后提出建议补充的图片。同时给出需补充图片的描述，以便根据描述制作图片。可以有多个补充图片的描述。\n"
    prompt = (
        f"Below are descriptions for a set of e-commerce images for the product {product_name}. "
        "For our target platform, a standard product image set should include at least 5 and up to 10 images. "
        "The mandatory images are: a main image (hero shot), at least one size/dimension image in both metric system and imperial system, and at least one lifestyle/usage scenario image. "
        "Optional images may include (but are not limited to): multiple use-case images, detailed close-ups, model shots, side views, and internal views. "
        "Please categorize the images based on their content and suggest any additional images needed to meet platform standards. "
        "Please also listed images that are not suitable for the platform or are low quality and give the reason."
        "Provide detailed descriptions for any recommended additional images to guide their creation. "
        "Multiple suggestions are welcome. Output text must be in Simplified Chinese and in markdown format.\nImage descriptions:\n"
    )
    # prompt += "\n".join([f"{file}: {desc}" for file, desc in descriptions])
    prompt += "\n".join(descriptions)
    
    model = DeepseekModel()
    try:
        response = model.chat(prompt)
        return response
    except Exception as e:
        print(f"Error analyzing descriptions: {e}")
        return None

def main():
    """
    主函数，用于运行产品图片审核工具。
    
    功能:
        - 提示用户输入产品名称和图片目录。
        - 检查图片尺寸是否符合要求。
        - 获取图片描述并分析是否需要补充图片。
        - 输出分析结果。
    """
    print("请确保输入的产品名称清晰且准确，以便系统更好地处理图片。")
    print("例如：'夏季新款男士T恤' 或 '智能蓝牙耳机黑色版'")
    while True:
        product_name = input("请输入产品名称：")
        if not product_name.strip():
            print("产品名称不能为空，请重新输入。")
            continue
        
        break

    while True:        
        image_dir = input("请输入图片所在目录：")
        if not os.path.isdir(image_dir):
            print("目录不存在或无法访问，请重新输入。")
            continue
        
        valid_images = check_image_dimensions(image_dir)
        if not valid_images:
            print("未找到符合尺寸要求的图片，请重新输入目录。")
            continue
        
        break
    
    descriptions, chinese_images = get_image_descriptions(product_name, image_dir)
    if not descriptions:
        print("未找到有效的图片描述。")
        return

    # 将描述信息保存为 Markdown 文件
    markdown_content = f"# {product_name}\n\n"
    for file, desc in descriptions:
        markdown_content += f"## {file}\n\n{desc}\n\n"
    md_file = os.path.join(image_dir, "图片描述.md")
    with open(md_file, 'w', encoding='utf-8') as f:
        f.write(markdown_content)
    print(f"图片描述已保存到文件: {md_file}")

    if chinese_images:
        print("\n以下图片描述中包含中文：")
        for image in chinese_images:
            print(f"- {image}")
    
    print("正在分析图片描述以提出补充建议...\n")
    analysis = analyze_descriptions(product_name, markdown_content)
    if analysis:
        # print("\n分析结果：")
        # print(analysis)
        if analysis.startswith("```"):
            analysis = '\n'.join(analysis.split("\n")[1:-1])
        show_markdown_popup(analysis, title=f"{product_name}图片分析结果")
        # 将分析结果保存到输入目录下的Markdown文件
        output_file = os.path.join(image_dir, f"{product_name}_analysis.md")
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(analysis)
        print(f"分析结果已保存到文件: {output_file}")
        

if __name__ == "__main__":
    main()